# 加载环境变量
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv
import os
import json
import requests

_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)

amap_key = "0638b10da58833b1c1efd6497d589dd1"


# 移除字典中的 None 值
def remove_none_values(data):
    if isinstance(data, dict):
        return {key: remove_none_values(value) for key, value in data.items() if value is not None}
    elif isinstance(data, list):
        return [remove_none_values(item) for item in data if item is not None]
    else:
        return data


# 打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
def print_json(data):
    if (isinstance(data, (list, dict))):
        print(json.dumps(data, indent=4, ensure_ascii=False))
    else:
        print(data)


# 对象序列化
def serialize_json(data):
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list, dict))):
        data = remove_none_values(data)
    return data


def get_completion(params, model="gpt-3.5-turbo"):
    print("=====GPT请求=====")
    messages = serialize_json(params)
    print_json(messages)

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
        seed=1024,  # 随机种子保持不变，temperature 和 prompt 不变的情况下，输出就会不变
        tool_choice="auto",  # 默认值，由 GPT 自主决定返回 function call 还是返回文字回复。也可以强制要求必须调用指定的函数，详见官方文档
        tools=[{
            "type": "function",
            "function": {
                "name": "get_location_coordinate",
                "description": "根据POI名称，获得POI的经纬度坐标",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "POI名称，必须是中文",
                        },
                        "city": {
                            "type": "string",
                            "description": "POI所在的城市名，必须是中文",
                        }
                    },
                    "required": ["location", "city"],
                }
            }
        },
            {
                "type": "function",
                "function": {
                    "name": "search_nearby_pois",
                    "description": "搜索给定坐标附近的poi",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "longitude": {
                                "type": "string",
                                "description": "中心点的经度",
                            },
                            "latitude": {
                                "type": "string",
                                "description": "中心点的纬度",
                            },
                            "keyword": {
                                "type": "string",
                                "description": "目标poi的关键字",
                            }
                        },
                        "required": ["longitude", "latitude", "keyword"],
                    }
                }
            }],
    )
    result = serialize_json(response.choices[0].message)
    print("=====GPT回复=====")
    print_json(result)
    return result


def get_location_coordinate(location, city):
    url = f"https://restapi.amap.com/v5/place/text?key={amap_key}&keywords={location}&region={city}"
    print("Call: get_location_coordinate", url)
    r = requests.get(url)
    response = r.json()
    if "pois" in response and response["pois"]:
        result = response["pois"][0]
        print("Return: get_location_coordinate", result)
        return result
    return "未找到该POI"


def search_nearby_pois(longitude, latitude, keyword):
    url = f"https://restapi.amap.com/v5/place/around?key={amap_key}&keywords={keyword}&location={longitude},{latitude}"
    print("Call: search_nearby_pois", url)
    r = requests.get(url)
    result = r.json()
    ans = ""
    if "pois" in result and result["pois"]:
        print("Return: search_nearby_pois", result["pois"])
        for i in range(min(3, len(result["pois"]))):
            name = result["pois"][i]["name"]
            address = result["pois"][i]["address"]
            distance = result["pois"][i]["distance"]
            ans += f"{name}\n{address}\n距离：{distance}米\n\n"
    else:
        ans = "没有找到附近的POI"
    return ans


prompt = "我想在北京五道口附近喝咖啡，给我推荐几个"
prompt = "我到北京出差，给我推荐三里屯的酒店，和海龙大厦附近的咖啡"
prompt = "我到北京出差，给我推荐三里屯的酒店，在中关村海龙大厦附近找一家星巴克，我要喝咖啡"

messages = [
    {"role": "system", "content": "你是一个地图通，你可以找到任何地址。"},
    {"role": "user", "content": prompt}
]
response = get_completion(messages)
messages.append(response)  # 把大模型的回复加入到对话中


while 'tool_calls' in response and response['tool_calls'] is not None:
    # 1106 版新模型支持一次返回多个函数调用请求，所以要考虑到这种情况
    for tool_call in response['tool_calls']:
        args = json.loads(tool_call['function']['arguments'])
        if tool_call['function']['name'] == "get_location_coordinate":
            result = get_location_coordinate(**args)
        elif tool_call['function']['name'] == "search_nearby_pois":
            result = search_nearby_pois(**args)

        messages.append({
            "tool_call_id": tool_call['id'],  # 用于标识函数调用的 ID
            "role": "tool",
            "name": tool_call['function']['name'],
            "content": str(result)  # 数值result 必须转成字符串
        })

    response = get_completion(messages)
    messages.append(response)  # 把大模型的回复加入到对话中

print("=====对话列表=====")
print_json(messages)

print("=====最终结果=====")
print(response['content'])
